import numpy
import os

if __name__ == '__main__':
    root = '/data/zw_1128/tracking_03'
    lidar_path = os.path.join(root, 'lidar')
    radar_path = os.path.join(root, 'radar')
    rgb_path = os.path.join(root, 'rgb')
    thermal_path = os.path.join(root, 'thermal')

    # read .csv file, first column is the new name, second column is the time stamp
    lidar_csv = numpy.loadtxt(os.path.join(root, 'lidar_timestamp.csv'), dtype=str, delimiter=',')
    radar_csv = numpy.loadtxt(os.path.join(root, 'radar_timestamp.csv'), dtype=str, delimiter=',')
    rgb_csv = numpy.loadtxt(os.path.join(root, 'image_timestamp.csv'), dtype=str, delimiter=',')
    thermal_csv = numpy.loadtxt(os.path.join(root, 'infared_timestamp.csv'), dtype=str, delimiter=',')

    # save path
    save_root = root + '_select'
    lidar_save_path = os.path.join(save_root, 'lidar')
    radar_save_path = os.path.join(save_root, 'radar')
    rgb_save_path = os.path.join(save_root, 'rgb')
    thermal_save_path = os.path.join(save_root, 'thermal')

    for path in [lidar_save_path, radar_save_path, rgb_save_path, thermal_save_path]:
        os.makedirs(path, exist_ok=True)

    # time stamp to be selected
    lidar_time = lidar_csv[:, 1].astype(float)
    radar_time = radar_csv[:, 1].astype(float)
    rgb_time = rgb_csv[:, 1].astype(float)
    thermal_time = thermal_csv[:, 1].astype(float)

    # search the closest time stamp in lidar, radar, rgb, thermal path
    lidar_files = sorted(os.listdir(lidar_path))
    lidar_timestamp_candidate = numpy.array([int(file.split('.')[0]) for file in lidar_files])
    # divide by 1e9 to convert to sec (float64)
    lidar_timestamp_candidate = lidar_timestamp_candidate / 1e9


    # copy the file to save path and rename, original name is the time stamp, new name is the first column in .csv file
    for i in range(len(lidar_time)):
        time = lidar_time[i]
        name = lidar_csv[i, 0].zfill(6)
        # find the closest time stamp
        idx = numpy.argmin(numpy.abs(lidar_timestamp_candidate - time))
        os.system('cp {} {}'.format(os.path.join(lidar_path, lidar_files[idx]), os.path.join(lidar_save_path, name + '.pcd')))




    radar_files = sorted(os.listdir(radar_path))
    radar_timestamp_candidate = numpy.array([int(file.split('.')[0]) for file in radar_files])
    radar_timestamp_candidate = radar_timestamp_candidate / 1e9
    for i in range(len(radar_time)):
        time = radar_time[i]
        name = radar_csv[i, 0].zfill(6)
        idx = numpy.argmin(numpy.abs(radar_timestamp_candidate - time))
        os.system('cp {} {}'.format(os.path.join(radar_path, radar_files[idx]), os.path.join(radar_save_path, name + '.pcd')))




    rgb_files = sorted(os.listdir(rgb_path))
    rgb_timestamp_candidate = numpy.array([int(file.split('.')[0]) for file in rgb_files])
    rgb_timestamp_candidate = rgb_timestamp_candidate / 1e9
    for i in range(len(rgb_time)):
        time = rgb_time[i]
        name = rgb_csv[i, 0].zfill(6)
        idx = numpy.argmin(numpy.abs(rgb_timestamp_candidate - time))
        os.system('cp {} {}'.format(os.path.join(rgb_path, rgb_files[idx]), os.path.join(rgb_save_path, name + '.png')))

    thermal_files = sorted(os.listdir(thermal_path))
    thermal_timestamp_candidate = numpy.array([int(file.split('.')[0]) for file in thermal_files])
    thermal_timestamp_candidate = thermal_timestamp_candidate / 1e9
    for i in range(len(thermal_time)):
        time = thermal_time[i]
        name = thermal_csv[i, 0].zfill(6)
        idx = numpy.argmin(numpy.abs(thermal_timestamp_candidate - time))
        os.system('cp {} {}'.format(os.path.join(thermal_path, thermal_files[idx]), os.path.join(thermal_save_path, name + '.png')))

